from __future__ import division
#Test Histogram equalization functions
from scipy.ndimage import *
import numpy as np
import pylab

def make_pdf(im,normed=True):
    """
    Calculate the normalized probability distribution function for image, im.  
    Return pdf as 1D array.
    """
    #make sure image is an array
    im = np.asarray(im,'f')
    h = histogram(im,0,254,256)
    if normed:
        pdf = h/h.sum()
        return pdf
    else:
        return h

def make_cdf(im):
    """
    Calculate the cumulative distribution function for image, im. 
    Return cdf as 1D array
    """
    #Get the pdf first 
    pdf = make_pdf(im)
    cdf = np.cumsum(pdf)
    return cdf

def histeq(im):
    """
    Perform a histogram equalization on image, im.  Return equalized image.
    """
    im = np.asarray(im,'int')
    cdf = make_cdf(im)
    im = 255*cdf[im]
    im = im.astype('int')
    return im

def plot(im):
    """
    Plot a before and after image grid
    """
    pdf1=make_pdf(im)
    cdf1=make_cdf(im)
    
    J = histeq(im)
    pdf2 = make_pdf(J)
    cdf2 = make_cdf(J)
    
    pylab.figure()
    pylab.subplot(221)
    pylab.imshow(im)
    pylab.subplot(222)
    pylab.plot(pdf1*1/pdf1.max())
    pylab.plot(cdf1)
    pylab.subplot(223)
    pylab.imshow(J)
    pylab.subplot(224)
    pylab.plot(pdf2*1/pdf2.max())
    pylab.plot(cdf2)
 
    